import os
import time
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, field
from threading import Lock

import pydicom

phase_dict = {
    "5%": 0,
    "15%": 1,
    "25%": 2,
    "35%": 3,
    "45%": 4,
    "55%": 5,
    "65%": 6,
    "75%": 7,
    "85%": 8,
    "95%": 9,
}


def turn_uid_dict_to_index_dict(series_dict, series_to_phase_dict):
    phase_index_dict = {}
    if len(series_to_phase_dict) == 0:
        for k in series_dict:
            phase_index_dict[0] = series_dict[k]
    else:
        for k in series_dict:
            phase_index = phase_dict[series_to_phase_dict[k]]
            phase_index_dict[phase_index] = series_dict[k]
    return phase_index_dict

    # 递归的方式读取文件夹下所有的dicom文件。


def get_dicom_files(directory):
    dicom_files = []
    for root, _, files in os.walk(directory):
        for file in files:
            if file.lower().endswith('.dcm'):
                dicom_files.append(os.path.join(root, file))
    return dicom_files


def load_give_fold_path_dicom(dicom_folder_path):
    # 使用并行处理加速文件读取
    from concurrent.futures import ThreadPoolExecutor

    def read_dicom(filepath):
        return pydicom.dcmread(filepath)

    files = [os.path.join(dicom_folder_path, f) for f in os.listdir(dicom_folder_path)]
    with ThreadPoolExecutor() as executor:
        temp = list(executor.map(read_dicom, files))

    slices = [t for t in temp if t.Modality == 'CT']
    slices.sort(key=lambda x: int(x.InstanceNumber))
    return slices


@dataclass
class Data4d:
    is_4d_series = False
    # 不添加field导致，在链接断开后，该对象没有释放。
    series_to_phase_dict: dict = field(default_factory=dict)


class DicomLoader:
    def __init__(self, dicom_folder_path):
        # series_dict 是一个 key是 seriesUID   value 是该series下的所有ds的数组。
        self.series_dict = {}
        self.data_4d = Data4d()

        # 是否是多序列。
        self.is_multi_series = False
        self.dicom_folder_path = dicom_folder_path

    def load_a_exist_fold(self, dicom_folder_path):
        self.dicom_folder_path = dicom_folder_path
        self.classify_dicom_files()
        self.turn_series_dict_to_order()
        self.check_is_multi_series()
        phase_index_dict = turn_uid_dict_to_index_dict(self.series_dict, self.data_4d.series_to_phase_dict)
        first_phase_index = -1
        if len(phase_index_dict) > 0:
            first_phase_index = list(phase_index_dict.keys())[0]
        return phase_index_dict, first_phase_index

    def load_a_dicom(self, dicom_path):
        ds = pydicom.dcmread(dicom_path)
        self.classify_a_dicom(ds)
        phase_index = 0
        if self.data_4d.is_4d_series:
            series_uid = ds.SeriesInstanceUID
            phase_index = phase_dict[self.data_4d.series_to_phase_dict[series_uid]]
        print('load a dicom,', dicom_path, phase_index)
        return ds, phase_index

    def check_is_multi_series(self):
        if len(self.series_dict) > 1:
            self.is_multi_series = True

    def classify_dicom_files(self):
        dicom_files = get_dicom_files(self.dicom_folder_path)
        start_time = time.time()
        for file_path in dicom_files:
            ds = pydicom.dcmread(file_path)
            self.classify_a_dicom(ds)

        end_time = time.time()
        print(f'classify_dicom_files,运行时间：{end_time - start_time}')

    def classify_a_dicom(self, ds):
        if hasattr(ds, 'SeriesInstanceUID') is False:
            raise ValueError("SeriesUID not exist")
        series_uid = ds.SeriesInstanceUID
        if series_uid in self.series_dict:
            self.series_dict[series_uid].append(ds)
        elif hasattr(ds, 'SeriesDescription') and (ds.SeriesDescription in phase_dict):
            self.data_4d.is_4d_series = True
            self.is_multi_series = True

            self.data_4d.series_to_phase_dict[series_uid] = ds.SeriesDescription
            self.series_dict[series_uid] = [ds]

        else:
            self.series_dict[series_uid] = [ds]

    def turn_series_dict_to_order(self):
        for series_id in self.series_dict:
            dicoms = self.series_dict[series_id]
            dicoms.sort(key=lambda x: int(x.InstanceNumber))
            self.series_dict[series_id] = dicoms
